#include <vector>
#include <cstring>
#include "vulkan_texture.h"
#include <vulkan/vulkan.h>
#include <rawfile/raw_file_manager.h>
#include <hilog/log.h>

#define STB_IMAGE_IMPLEMENTATION
#include "stb/stb_image.h"

// 修复1: 避免与系统宏冲突，使用不同的名称
static const char* VULKAN_TEXTURE_TAG = "VulkanTexture";

// 修复2: 定义日志域
#ifndef LOG_DOMAIN
#define LOG_DOMAIN 0x3200  // 使用自定义域ID
#endif

VkResult createBuffer(VkDevice device, VkPhysicalDevice physicalDevice, VkDeviceSize size,
                     VkBufferUsageFlags usage, VkMemoryPropertyFlags properties,
                     VkBuffer* buffer, VkDeviceMemory* bufferMemory) {
    VkBufferCreateInfo bufferInfo{};
    bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
    bufferInfo.size = size;
    bufferInfo.usage = usage;
    bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    
    VkResult result = vkCreateBuffer(device, &bufferInfo, nullptr, buffer);
    if (result != VK_SUCCESS) {
        return result;
    }
    
    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(device, *buffer, &memRequirements);
    
    VkMemoryAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocInfo.allocationSize = memRequirements.size;
    allocInfo.memoryTypeIndex = findMemoryType(physicalDevice, memRequirements.memoryTypeBits, properties);
    
    result = vkAllocateMemory(device, &allocInfo, nullptr, bufferMemory);
    if (result != VK_SUCCESS) {
        vkDestroyBuffer(device, *buffer, nullptr);
        return result;
    }
    
    vkBindBufferMemory(device, *buffer, *bufferMemory, 0);
    return VK_SUCCESS;
}

uint32_t findMemoryType(VkPhysicalDevice physicalDevice, uint32_t typeFilter, 
                       VkMemoryPropertyFlags properties) {
    VkPhysicalDeviceMemoryProperties memProperties;
    vkGetPhysicalDeviceMemoryProperties(physicalDevice, &memProperties);
    
    for (uint32_t i = 0; i < memProperties.memoryTypeCount; i++) {
        if ((typeFilter & (1 << i)) && 
            (memProperties.memoryTypes[i].propertyFlags & properties) == properties) {
            return i;
        }
    }
    
    // 如果找不到合适的内存类型，抛出异常或返回无效值
    OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, VULKAN_TEXTURE_TAG, "Failed to find suitable memory type!");
    return UINT32_MAX; // 表示失败
}

VkCommandBuffer beginSingleTimeCommands(VkDevice device, VkCommandPool commandPool) {
    VkCommandBufferAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocInfo.commandPool = commandPool;
    allocInfo.commandBufferCount = 1;
    
    VkCommandBuffer commandBuffer;
    vkAllocateCommandBuffers(device, &allocInfo, &commandBuffer);
    
    VkCommandBufferBeginInfo beginInfo{};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
    
    vkBeginCommandBuffer(commandBuffer, &beginInfo);
    return commandBuffer;
}

void endSingleTimeCommands(VkDevice device, VkCommandPool commandPool, 
                          VkQueue queue, VkCommandBuffer commandBuffer) {
    vkEndCommandBuffer(commandBuffer);
    
    VkSubmitInfo submitInfo{};
    submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers = &commandBuffer;
    
    vkQueueSubmit(queue, 1, &submitInfo, VK_NULL_HANDLE);
    vkQueueWaitIdle(queue);
    
    vkFreeCommandBuffers(device, commandPool, 1, &commandBuffer);
}

VkResult transitionImageLayout(VkDevice device, VkCommandPool commandPool, VkQueue queue,
                              VkImage image, VkFormat format, VkImageLayout oldLayout,
                              VkImageLayout newLayout) {
    VkCommandBuffer commandBuffer = beginSingleTimeCommands(device, commandPool);
    
    VkImageMemoryBarrier barrier{};
    barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
    barrier.oldLayout = oldLayout;
    barrier.newLayout = newLayout;
    barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.image = image;
    barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    barrier.subresourceRange.baseMipLevel = 0;
    barrier.subresourceRange.levelCount = 1;
    barrier.subresourceRange.baseArrayLayer = 0;
    barrier.subresourceRange.layerCount = 1;
    
    VkPipelineStageFlags sourceStage;
    VkPipelineStageFlags destinationStage;
    
    if (oldLayout == VK_IMAGE_LAYOUT_UNDEFINED && newLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL) {
        barrier.srcAccessMask = 0;
        barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
        
        sourceStage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
        destinationStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
    } else if (oldLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL && newLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL) {
        barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
        barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
        
        sourceStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
        destinationStage = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
    } else {
        OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, VULKAN_TEXTURE_TAG, "Unsupported layout transition!");
        endSingleTimeCommands(device, commandPool, queue, commandBuffer);
        return VK_ERROR_FEATURE_NOT_PRESENT;
    }
    
    vkCmdPipelineBarrier(commandBuffer, sourceStage, destinationStage, 0, 0, nullptr, 0, nullptr, 1, &barrier);
    
    endSingleTimeCommands(device, commandPool, queue, commandBuffer);
    return VK_SUCCESS;
}

VkResult copyBufferToImage(VkDevice device, VkCommandPool commandPool, VkQueue queue,
                          VkBuffer buffer, VkImage image, uint32_t width, uint32_t height) {
    VkCommandBuffer commandBuffer = beginSingleTimeCommands(device, commandPool);
    
    VkBufferImageCopy region{};
    region.bufferOffset = 0;
    region.bufferRowLength = 0;
    region.bufferImageHeight = 0;
    region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    region.imageSubresource.mipLevel = 0;
    region.imageSubresource.baseArrayLayer = 0;
    region.imageSubresource.layerCount = 1;
    region.imageOffset = {0, 0, 0};
    region.imageExtent = {width, height, 1};
    
    vkCmdCopyBufferToImage(commandBuffer, buffer, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &region);
    
    endSingleTimeCommands(device, commandPool, queue, commandBuffer);
    return VK_SUCCESS;
}

VkResult createTexture(VkDevice device, VkPhysicalDevice physicalDevice, 
                      VkCommandPool commandPool, VkQueue queue,  // 修复3: 添加缺少的参数
                      const char* imagePath, VkImage* textureImage, 
                      VkDeviceMemory* textureImageMemory) {
    // 从rawfile加载图像数据
    int texWidth, texHeight, texChannels;
    stbi_uc* pixels = loadImageFromRawFile(imagePath, &texWidth, &texHeight, &texChannels);
    
    if (!pixels) {
        // 修复2: 使用正确的日志函数调用格式
        OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, VULKAN_TEXTURE_TAG, "Failed to load texture image: %s", imagePath);
        return VK_ERROR_INITIALIZATION_FAILED;
    }
    
    VkDeviceSize imageSize = texWidth * texHeight * 4; // RGBA
    
    // 创建暂存缓冲区
    VkBuffer stagingBuffer;
    VkDeviceMemory stagingBufferMemory;
    VkResult result = createBuffer(device, physicalDevice, imageSize, 
                VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
                &stagingBuffer, &stagingBufferMemory);
    
    if (result != VK_SUCCESS) {
        stbi_image_free(pixels);
        return result;
    }
    
    // 将像素数据复制到暂存缓冲区
    void* data;
    result = vkMapMemory(device, stagingBufferMemory, 0, imageSize, 0, &data);
    if (result != VK_SUCCESS) {
        stbi_image_free(pixels);
        vkDestroyBuffer(device, stagingBuffer, nullptr);
        vkFreeMemory(device, stagingBufferMemory, nullptr);
        return result;
    }
    
    memcpy(data, pixels, static_cast<size_t>(imageSize));
    vkUnmapMemory(device, stagingBufferMemory);
    
    stbi_image_free(pixels);
    
    // 创建图像
    VkImageCreateInfo imageInfo{};
    imageInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
    imageInfo.imageType = VK_IMAGE_TYPE_2D;
    imageInfo.extent.width = static_cast<uint32_t>(texWidth);
    imageInfo.extent.height = static_cast<uint32_t>(texHeight);
    imageInfo.extent.depth = 1;
    imageInfo.mipLevels = 1;
    imageInfo.arrayLayers = 1;
    imageInfo.format = VK_FORMAT_R8G8B8A8_SRGB;
    imageInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
    imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
    imageInfo.usage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT;
    imageInfo.samples = VK_SAMPLE_COUNT_1_BIT;
    imageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    
    result = vkCreateImage(device, &imageInfo, nullptr, textureImage);
    if (result != VK_SUCCESS) {
        vkDestroyBuffer(device, stagingBuffer, nullptr);
        vkFreeMemory(device, stagingBufferMemory, nullptr);
        return result;
    }
    
    // 分配图像内存
    VkMemoryRequirements memRequirements;
    vkGetImageMemoryRequirements(device, *textureImage, &memRequirements);
    
    VkMemoryAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocInfo.allocationSize = memRequirements.size;
    allocInfo.memoryTypeIndex = findMemoryType(physicalDevice, memRequirements.memoryTypeBits,
                                              VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
    
    result = vkAllocateMemory(device, &allocInfo, nullptr, textureImageMemory);
    if (result != VK_SUCCESS) {
        vkDestroyImage(device, *textureImage, nullptr);
        vkDestroyBuffer(device, stagingBuffer, nullptr);
        vkFreeMemory(device, stagingBufferMemory, nullptr);
        return result;
    }
    
    vkBindImageMemory(device, *textureImage, *textureImageMemory, 0);
    
    // 转换图像布局并复制数据
    result = transitionImageLayout(device, commandPool, queue, *textureImage, 
                         VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_UNDEFINED, 
                         VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
    if (result != VK_SUCCESS) {
        // 清理资源
        vkDestroyImage(device, *textureImage, nullptr);
        vkFreeMemory(device, *textureImageMemory, nullptr);
        vkDestroyBuffer(device, stagingBuffer, nullptr);
        vkFreeMemory(device, stagingBufferMemory, nullptr);
        return result;
    }
    
    result = copyBufferToImage(device, commandPool, queue, stagingBuffer, *textureImage, 
                     static_cast<uint32_t>(texWidth), static_cast<uint32_t>(texHeight));
    if (result != VK_SUCCESS) {
        // 清理资源
        vkDestroyImage(device, *textureImage, nullptr);
        vkFreeMemory(device, *textureImageMemory, nullptr);
        vkDestroyBuffer(device, stagingBuffer, nullptr);
        vkFreeMemory(device, stagingBufferMemory, nullptr);
        return result;
    }
    
    result = transitionImageLayout(device, commandPool, queue, *textureImage, 
                         VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 
                         VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
    if (result != VK_SUCCESS) {
        // 清理资源
        vkDestroyImage(device, *textureImage, nullptr);
        vkFreeMemory(device, *textureImageMemory, nullptr);
    }
    
    // 清理暂存缓冲区
    vkDestroyBuffer(device, stagingBuffer, nullptr);
    vkFreeMemory(device, stagingBufferMemory, nullptr);
    
    return result;
}

stbi_uc* loadImageFromRawFile(const char* imagePath, int* width, int* height, int* channels) {
    // 获取资源管理器
    NativeResourceManager* nativeResourceManager = OH_ResourceManager_InitNativeResourceManager(nullptr, nullptr);
    if (!nativeResourceManager) {
        OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, VULKAN_TEXTURE_TAG, "Failed to init resource manager");
        return nullptr;
    }
    
    // 打开rawfile
    RawFile* rawFile = OH_ResourceManager_OpenRawFile(nativeResourceManager, imagePath);
    if (!rawFile) {
        OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, VULKAN_TEXTURE_TAG, "Failed to open raw file: %s", imagePath);
        OH_ResourceManager_ReleaseNativeResourceManager(nativeResourceManager);
        return nullptr;
    }
    
    // 获取文件大小
    long fileSize = OH_ResourceManager_GetRawFileSize(rawFile);
    if (fileSize <= 0) {
        OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, VULKAN_TEXTURE_TAG, "Invalid file size: %ld", fileSize);
        OH_ResourceManager_CloseRawFile(rawFile);
        OH_ResourceManager_ReleaseNativeResourceManager(nativeResourceManager);
        return nullptr;
    }
    
    // 读取文件数据
    std::vector<unsigned char> buffer(fileSize);
    long bytesRead = OH_ResourceManager_ReadRawFile(rawFile, buffer.data(), fileSize);
    if (bytesRead != fileSize) {
        OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, VULKAN_TEXTURE_TAG, "Failed to read complete file");
        OH_ResourceManager_CloseRawFile(rawFile);
        OH_ResourceManager_ReleaseNativeResourceManager(nativeResourceManager);
        return nullptr;
    }
    
    // 使用stb_image解码
    stbi_uc* pixels = stbi_load_from_memory(buffer.data(), static_cast<int>(fileSize), 
                                           width, height, channels, STBI_rgb_alpha);
    
    if (!pixels) {
        OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, VULKAN_TEXTURE_TAG, "Failed to decode image: %s", stbi_failure_reason());
    }
    
    // 清理资源
    OH_ResourceManager_CloseRawFile(rawFile);
    OH_ResourceManager_ReleaseNativeResourceManager(nativeResourceManager);
    
    return pixels;
}


